#include <linux/init.h>
#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/netlink.h>
#include <net/netlink.h>
#include <net/net_namespace.h>
#include <linux/skbuff.h>

#define NETLINK_USER 31  /* 自定义协议号，实际使用时需避免冲突 */

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Your Name");
MODULE_DESCRIPTION("Simple Netlink Kernel Module");
MODULE_VERSION("0.1");

/* 用于保存内核侧创建的 netlink socket */
static struct sock *nlk_sock = NULL;

/* 收到消息后的处理函数 */
static void nlk_recv_msg(struct sk_buff *skb)
{
    struct nlmsghdr *nlh;
    struct sk_buff  *skb_out;
    int    pid;
    int    msg_size;
    char   *msg = "Hello from kernel";
    int    res;

    /* 从 sk_buff 提取出 netlink 消息头 */
    nlh = (struct nlmsghdr *)skb->data;
    pr_info("Kernel module received msg payload: %s\n", (char *)nlmsg_data(nlh));

    /* 获取发送者的 PID（即用户态进程的 netlink port ID） */
    pid = nlh->nlmsg_pid;

    /* 构造要发送回用户态的消息 */
    msg_size = strlen(msg);
    /* 为要发送的消息申请一个新的 sk_buff */
    skb_out = nlmsg_new(msg_size, GFP_KERNEL);
    if (!skb_out) {
        pr_err("Failed to allocate new skb\n");
        return;
    }

    /* 在 skb_out 中填充 netlink 消息头和消息数据 */
    nlh = nlmsg_put(skb_out, 0, 0, NLMSG_DONE, msg_size, 0);
    if (!nlh) {
        pr_err("nlmsg_put failed\n");
        kfree_skb(skb_out);
        return;
    }
    memcpy(nlmsg_data(nlh), msg, msg_size);

    /* 发送消息给对应 pid 的用户进程 */
    res = nlmsg_unicast(nlk_sock, skb_out, pid);
    if (res < 0) {
        pr_err("Error while sending back to user\n");
    }
}

/* 模块初始化函数 */
static int __init nlk_init(void)
{
    struct netlink_kernel_cfg cfg = {
        .input = nlk_recv_msg,  /* 收到消息时的回调函数 */
    };

    /* 创建一个 Netlink 内核 socket (PF_NETLINK) 并指定回调 */
    nlk_sock = netlink_kernel_create(&init_net, NETLINK_USER, &cfg);
    if (!nlk_sock) {
        pr_err("Error creating socket.\n");
        return -10;
    }

    pr_info("Netlink Kernel Module inserted.\n");
    return 0;
}

/* 模块卸载函数 */
static void __exit nlk_exit(void)
{
    netlink_kernel_release(nlk_sock);
    pr_info("Netlink Kernel Module removed.\n");
}

module_init(nlk_init);
module_exit(nlk_exit);
